package liming.tool.encryption;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.Serializable;
import java.security.*;

/**
 * 混合加密工具类，结合RSA非对称加密与AES对称加密的优势：
 * - 使用RSA安全传输AES密钥
 * - 使用AES-GCM高效加密业务数据
 * 包含服务端和客户端处理逻辑的嵌套类实现
 */
public class HybridEncryption {
    //===================== 算法参数常量 ========================//
    /** AES密钥默认长度（256位） */
    protected static final int AES_KEY_SIZE = 256;
    /** RSA密钥默认长度（2048位） */
    protected static final int RSA_KEY_SIZE = 2048;
    /** GCM随机数默认长度（12字节，推荐值） */
    protected static final int GCM_NONCE_LENGTH = 12;
    /** GCM认证标签默认长度（128位，提供强完整性校验） */
    protected static final int GCM_TAG_LENGTH = 128;

    //===================== 核心加密方法 ========================//

    /**
     * 生成AES对称密钥
     * @param aesKeySize 密钥长度（单位：bit，建议256）
     * @return 符合指定长度的AES密钥
     * @throws NoSuchAlgorithmException 当JCE不支持AES算法时抛出
     */
    private static SecretKey generateAESKey(int aesKeySize) throws NoSuchAlgorithmException {
        KeyGenerator keyGen = KeyGenerator.getInstance("AES");
        keyGen.init(aesKeySize);
        return keyGen.generateKey();
    }

    /**
     * RSA公钥加密数据（使用OAEP填充方案）
     * @param publicKey RSA公钥对象
     * @param data 待加密的原始数据
     * @return 经过RSA加密的字节数组
     * @throws Exception 加密过程异常时抛出（如密钥不匹配、数据过长等）
     */
    private static byte[] rsaEncrypt(PublicKey publicKey, byte[] data) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        return cipher.doFinal(data);
    }

    /**
     * RSA私钥解密数据（需与加密填充方案一致）
     * @param privateKey RSA私钥对象
     * @param encryptedData 已加密的字节数据
     * @return 解密后的原始字节数据
     * @throws Exception 解密失败时抛出（如密钥错误、数据篡改等）
     */
    private static byte[] rsaDecrypt(PrivateKey privateKey, byte[] encryptedData) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        return cipher.doFinal(encryptedData);
    }

    /**
     * AES-GCM模式加密（自动生成并附加Nonce）
     * @param gcmNonceLength Nonce字节长度（推荐12字节）
     * @param gcmTagLength 认证标签位长度（推荐128位）
     * @param secretKey 有效的AES密钥
     * @param data 待加密的明文数据
     * @return 包含Nonce+密文的结构化字节数组
     * @throws Exception 加密过程异常时抛出（如密钥无效、参数错误等）
     */
    protected static byte[] aesEncrypt(int gcmNonceLength, int gcmTagLength, SecretKey secretKey, byte[] data) throws Exception {
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        // 生成密码学安全的随机Nonce
        byte[] nonce = new byte[gcmNonceLength];
        SecureRandom random = new SecureRandom();
        random.nextBytes(nonce);

        // 初始化GCM参数
        GCMParameterSpec spec = new GCMParameterSpec(gcmTagLength, nonce);
        cipher.init(Cipher.ENCRYPT_MODE, secretKey, spec);
        byte[] encryptedData = cipher.doFinal(data);

        // 拼接Nonce和密文（兼容格式）
        byte[] encryptedDataWithNonce = new byte[nonce.length + encryptedData.length];
        System.arraycopy(nonce, 0, encryptedDataWithNonce, 0, nonce.length);
        System.arraycopy(encryptedData, 0, encryptedDataWithNonce, nonce.length, encryptedData.length);
        return encryptedDataWithNonce;
    }

    /**
     * AES-GCM模式解密（解析Nonce+密文结构）
     * @param gcmNonceLength Nonce字节长度（需与加密时一致）
     * @param gcmTagLength 认证标签位长度（需与加密时一致）
     * @param secretKey 解密用的AES密钥
     * @param encryptedDataWithNonce 包含Nonce的加密数据
     * @return 解密后的原始字节数据
     * @throws Exception 解密失败时抛出（如密钥错误、数据篡改等）
     */
    protected static byte[] aesDecrypt(int gcmNonceLength, int gcmTagLength, SecretKey secretKey, byte[] encryptedDataWithNonce) throws Exception {
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");

        // 分离Nonce和密文部分
        byte[] nonce = new byte[gcmNonceLength];
        System.arraycopy(encryptedDataWithNonce, 0, nonce, 0, nonce.length);

        byte[] encryptedData = new byte[encryptedDataWithNonce.length - gcmNonceLength];
        System.arraycopy(encryptedDataWithNonce, gcmNonceLength, encryptedData, 0, encryptedData.length);

        // 重建GCM参数进行解密
        GCMParameterSpec spec = new GCMParameterSpec(gcmTagLength, nonce);
        cipher.init(Cipher.DECRYPT_MODE, secretKey, spec);
        return cipher.doFinal(encryptedData);
    }

    //===================== 基础功能类 ========================//

    /**
     * 混合加密基础类，封装通用加密参数和操作
     * @implNote 实现Serializable支持序列化传输
     */
    public static class HybridEncryptionBase implements Serializable {
        /** 当前实例持有的AES密钥 */
        protected SecretKey aesKey;
        /** GCM随机数长度（默认12字节） */
        private int gcmNonceLength = GCM_NONCE_LENGTH;
        /** GCM认证标签长度（默认128位） */
        private int gcmTagLength = GCM_TAG_LENGTH;

        /**
         * 设置AES对称密钥
         * @param aesKey 有效的AES密钥对象
         */
        public SecretKey setAesKey(SecretKey aesKey) {
            return this.aesKey = aesKey;
        }

        /**
         * 配置GCM模式参数
         * @param gcmNonceLength Nonce字节长度（建议12）
         * @param gcmTagLength 认证标签位长度（建议128）
         */
        public void setGcmLength(int gcmNonceLength, int gcmTagLength) {
            this.gcmNonceLength = gcmNonceLength;
            this.gcmTagLength = gcmTagLength;
        }

        /**
         * 执行数据解密
         * @param encryptedData 包含Nonce的加密数据
         * @return 解密后的原始数据
         * @throws Exception 解密失败时抛出（含详细错误信息）
         */
        public byte[] decrypt(byte[] encryptedData) throws Exception {
            if (aesKey == null)
                throw new RuntimeException("AES密钥未初始化");
            return aesDecrypt(gcmNonceLength, gcmTagLength, aesKey, encryptedData);
        }

        /**
         * 执行数据加密
         * @param data 待加密的明文数据
         * @return 包含Nonce的加密数据包
         * @throws Exception 加密失败时抛出（含详细错误信息）
         */
        public byte[] encrypt(byte[] data) throws Exception {
            if (aesKey == null)
                throw new RuntimeException("AES密钥未初始化");
            return aesEncrypt(gcmNonceLength, gcmTagLength, aesKey, data);
        }
    }

    //===================== 服务端实现 ========================//

    /**
     * 服务端处理模块，主要职责：
     * - 生成/管理RSA密钥对
     * - 解密客户端传输的AES密钥
     * @implNote 继承基础加密功能，支持序列化
     */
    public static class Server extends HybridEncryptionBase implements Serializable {
        /** RSA密钥对（包含公钥和私钥） */
        private KeyPair rsaKeyPair;

        /**
         * 生成默认长度的RSA密钥对（2048位）
         * @throws NoSuchAlgorithmException 当环境不支持RSA算法时抛出
         */
        public PublicKey newRsaKeyPair() throws NoSuchAlgorithmException {
            return newRsaKeyPair(RSA_KEY_SIZE);
        }

        /**
         * 生成指定长度的RSA密钥对
         * @param rsaKeySize 密钥长度（单位：bit，建议2048+）
         * @throws NoSuchAlgorithmException 当环境不支持RSA算法时抛出
         */
        public PublicKey newRsaKeyPair(int rsaKeySize) throws NoSuchAlgorithmException {
            KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
            keyGen.initialize(rsaKeySize);
            setRsaKeyPair(keyGen.generateKeyPair());
            return getPublicKey();
        }

        /**
         * 获取完整的RSA密钥对
         * @return 包含公钥和私钥的密钥对
         */
        public KeyPair getRsaKeyPair() {
            return rsaKeyPair;
        }

        /**
         * 设置RSA密钥对（用于密钥轮换场景）
         * @param rsaKeyPair 预先生成的有效密钥对
         */
        public void setRsaKeyPair(KeyPair rsaKeyPair) {
            this.rsaKeyPair = rsaKeyPair;
        }

        /**
         * 获取RSA公钥（需分发给客户端）
         * @return 公钥对象
         */
        public PublicKey getPublicKey() {
            return rsaKeyPair.getPublic();
        }

        /**
         * 解密客户端传来的AES密钥
         * @param encryptedAesKey RSA加密后的AES密钥数据
         * @throws Exception 解密失败时抛出（如私钥不匹配）
         */
        public void decryptAesKey(byte[] encryptedAesKey) throws Exception {
            // 使用RSA私钥解密后重建AES密钥
            byte[] rawKey = rsaDecrypt(rsaKeyPair.getPrivate(), encryptedAesKey);
            setAesKey(new SecretKeySpec(rawKey, "AES"));
        }
    }

    //===================== 客户端实现 ========================//

    /**
     * 客户端处理模块，主要职责：
     * - 生成AES会话密钥
     * - 加密AES密钥供服务端解密
     * @implNote 继承基础加密功能，支持序列化
     */
    public static class Client extends HybridEncryptionBase implements Serializable {
        /** 服务端提供的RSA公钥 */
        private PublicKey publicKey;

        /**
         * 设置服务端公钥（需在加密前调用）
         * @param publicKey 有效的RSA公钥
         */
        public PublicKey setPublic(PublicKey publicKey) {
            return this.publicKey = publicKey;
        }

        /**
         * 生成默认长度的AES密钥（256位）
         * @throws NoSuchAlgorithmException 当环境不支持AES时抛出
         */
        public SecretKey newAesKey() throws NoSuchAlgorithmException {
            return newAesKey(AES_KEY_SIZE);
        }

        /**
         * 生成指定长度的AES密钥
         * @param aesKeySize 密钥长度（单位：bit，建议256）
         * @throws NoSuchAlgorithmException 当环境不支持AES时抛出
         */
        public SecretKey newAesKey(int aesKeySize) throws NoSuchAlgorithmException {
            return setAesKey(generateAESKey(aesKeySize));
        }

        /**
         * 获取当前AES密钥（调试用）
         * @return AES密钥对象
         */
        public SecretKey getAesKey() {
            return aesKey;
        }

        /**
         * 加密AES密钥供服务端解密
         * @return 经RSA公钥加密的AES密钥数据
         * @throws Exception 加密失败时抛出（如公钥无效）
         */
        public byte[] encryptAesKey() throws Exception {
            if (publicKey == null)
                throw new RuntimeException("未配置服务端公钥");
            if (aesKey == null)
                throw new RuntimeException("未生成AES会话密钥");
            return rsaEncrypt(publicKey, aesKey.getEncoded());
        }
    }
}